from sklearn.model_selection import train_test_split
from sklearn.model_selection import RandomizedSearchCV, GridSearchCV
from sklearn.metrics import r2_score,accuracy_score
from sklearn.ensemble import RandomForestRegressor,RandomForestClassifier
import xgboost as xgb
from scipy.stats import randint
import numpy as np
import joblib
from logic.config import Config
import os

class MLModels:
    def __init__(self, task_name):
        self.task_name = task_name
        self.model_name = Config.get_config_path() + '/models/' + self.task_name + '.model'

    def RF(self,X,y):
        X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
        param_distributions = {'n_estimators': randint(1, 5),
                               'max_depth': randint(5, 10)}

        model = RandomForestClassifier(random_state=0)
        #model = RandomForestRegressor(random_state=0)

        # now create a searchCV object and fit it to the data
        search = RandomizedSearchCV(estimator=model,
                                    n_iter=5,
                                    param_distributions=param_distributions,
                                    random_state=0)
        search.fit(X_train, y_train)
        print(search.best_params_)
        print('测试集R2:', search.score(X_test, y_test))
        print('训练集R2:', search.score(X_train, y_train))
        #print('accuarcy:', accuracy_score(y_test, y_pred))
        #print('r2_score:', r2_score(y_test, y_pred))  # 这个顺序不能反。
        return search

    def is_model_exist(self):
        return os.path.exists(self.model_name)

    def dump_model(self, model):
        joblib.dump(model, self.model_name)

    def train(self, X, y):
        return self.RF(X, y)

    def predict(self, X, y):
        if self.is_model_exist():
            model = joblib.load(self.model_name)
        else:
            model = self.train(X, y)
            self.dump_model(model)

        y_pred = model.predict(X)
        return y_pred

if __name__ == '__main__':
    from logic.global_objs import D
    from engine.data.data_handler import DataHandler

    fields, names = DataHandler().get_kbar_fields_names()

    '''

    fields = ['Return($close,20)',
              'Return($close,10)',
              'Return($close,5)',

              ]

    names = fields
    '''

    all_names = names.copy()
    all_fields = fields.copy()

    all_fields.append('Ref($close,-5)/$close -1')
    all_names.append('label')

    all_fields.append('QCut($label,5)')
    all_names.append('label_c')

    df_all = D.load(['000300.SH', '000905.SH', '399006.SZ'], start_time='20100101', fields=all_fields, names=all_names)
    print(df_all)

    model = MLModels('我的任务')
    X, y = df_all[names], df_all['label_c']
    pred = model.predict(X, y)
    print(pred)

